#include "FlvWriter.h"

static uint8_t SCRIPT_TAG[FLV_SCRIPT_TAG_LENGTH] = {
    0x12, 0x00, 0x00, 0xB8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,//TagHeader 长度11
    0x02, 0x00, 0x0A, 0x6F, 0x6E, 0x4D, 0x65, 0x74, 0x61, 0x44, 0x61, 0x74, 0x61, 0x08, 0x00, 0x00, 0x00, 0x08, //onMetaData 数组长度为8
    0x00, 0x08, 0x64, 0x75, 0x72, 0x61, 0x74, 0x69, 0x6F, 0x6E, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,//duration
    0x00, 0x05, 0x77, 0x69, 0x64, 0x74, 0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,//width
    0x00, 0x06, 0x68, 0x65, 0x69, 0x67, 0x68, 0x74, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,//height
    0x00, 0x0D, 0x76, 0x69, 0x64, 0x65, 0x6F, 0x64, 0x61, 0x74, 0x61, 0x72, 0x61, 0x74, 0x65, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,//videodatarate
    0x00, 0x09, 0x66, 0x72, 0x61, 0x6D, 0x65, 0x72, 0x61, 0x74, 0x65, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,//framerate
    0x00, 0x0C, 0x76, 0x69, 0x64, 0x65, 0x6F, 0x63, 0x6F, 0x64, 0x65, 0x63, 0x69, 0x64, 0x00, 0x40, 0x1C, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,//videocodecid
    0x00, 0x07, 0x65, 0x6E, 0x63, 0x6F, 0x64, 0x65, 0x72, 0x02, 0x00, 0x0D, 0x4C, 0x61, 0x76, 0x66, 0x35, 0x38, 0x2E, 0x32, 0x30, 0x2E, 0x31, 0x30, 0x30,//encoder
    0x00, 0x08, 0x66, 0x69, 0x6C, 0x65, 0x73, 0x69, 0x7A, 0x65, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,//filesize
    0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0xC3//结束码和TAG长度
};
static uint8_t flvPre[FLV_HEADER_LENGTH] = {0x46, 0x4c, 0x56, 0x01, 0x01, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x00};
static uint8_t AVCSeq[400] = {
    0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x17, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0xFF, 0xE1, 0x00, 0x00
};
static uint8_t FLV_TAG_HEADER[FLV_TAG_HEADER_LENGTH] = {0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
static uint8_t PFRAME_PREFIX[PFRAME_PREFIX_LENGTH] = {0x27, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x61};
static std::string CHROMA_FORMAT[4] = {"单色", "YUV420", "YUV422", "YUV420"};

bool FlvWriter::SPSHandler()
{
    if(h264SPSLength != 0) return true;

    H264SPS[0] = 0x67;
    memcpy(H264SPS + 1, recBuffer + sendNode->startIndex, curNaluLength);
    h264SPSLength = curNaluLength + 1;

    if(h264PPSLength != 0)
    {
        if(!AfterGetSPSAndPPS()) return false;
    }
    return true;
}

bool FlvWriter::PPSHandler()
{
    if(h264PPSLength != 0) return true;

    H264PPS[0] = 0x68;
    memcpy(H264PPS + 1, recBuffer + sendNode->startIndex, curNaluLength);
    h264PPSLength = curNaluLength + 1;

    if(h264SPSLength != 0)
    {
        if(!AfterGetSPSAndPPS()) return false;
    }
    return true;
}

bool FlvWriter::IDRHandler()
{
    if(h264SPSLength == 0 || h264PPSLength == 0) return true;

    if(sendNode->packageType == 0 || sendNode->packageType == 1)
    {
        //处理帧长度
        int frameLength = htonl(IDRPrefixLength + curNaluLength);
        memcpy(FLV_TAG_HEADER, &frameLength, 4);
        FLV_TAG_HEADER[0] = 0x09;

        //处理时间戳
        int timestamp = htonl(sendNode->timestamp / 100);
        uint8_t temp[4];
        memcpy(temp, &timestamp, 4);
        memcpy(FLV_TAG_HEADER + 4, temp + 1, 3);
        FLV_TAG_HEADER[7] = temp[0];

        //处理IDRPrefixLength
        int idrLength = htonl(curNaluLength + 1);
        memcpy(IDRPrefix + 13 + h264SPSLength + h264PPSLength, &idrLength, 4);

        if(!Tools::WriteFile(file, clientFd, FLV_TAG_HEADER, FLV_TAG_HEADER_LENGTH, "写入FlvHeader时发生错误。")) return false;
        if(!Tools::WriteFile(file, clientFd, IDRPrefix, IDRPrefixLength, "写入IDRPrefix时发生错误。")) return false;   
    }
    
    int payloadLength = sendNode->endIndex - sendNode->startIndex;
    if(payloadLength > 0)
    {
        if(!Tools::WriteFile(file, clientFd, recBuffer + sendNode->startIndex, payloadLength, "写入nalu载荷时发生错误。")) return false;
    }
    else
    {
        payloadLength = REC_BUFFER_LENGTH_FLV - sendNode->startIndex;
        if(!Tools::WriteFile(file, clientFd, recBuffer + sendNode->startIndex, payloadLength, "写入nalu载荷时发生错误。")) return false;

        payloadLength = sendNode->endIndex;
        if(!Tools::WriteFile(file, clientFd, recBuffer, payloadLength, "写入nalu载荷时发生错误。")) return false;
    }

    if(sendNode->packageType == 0 || sendNode->packageType == 2)
    {
        int previousframeLength = htonl(IDRPrefixLength + curNaluLength + FLV_TAG_HEADER_LENGTH);
        if(!Tools::WriteFile(file, clientFd, &previousframeLength, 4, "写入上一帧长度时发生错误。")) return false;
    }

    return true;
}

bool FlvWriter::PFrameHandler()
{
    if(h264SPSLength == 0 || h264PPSLength == 0) return true;

    if(sendNode->packageType == 0 || sendNode->packageType == 1)
    {
        //处理帧长度
        int frameLength = htonl(PFRAME_PREFIX_LENGTH + curNaluLength);
        memcpy(FLV_TAG_HEADER, &frameLength, 4);
        FLV_TAG_HEADER[0] = 0x09;

        //处理时间戳
        int timestamp = htonl(sendNode->timestamp / 100);
        uint8_t temp[4];
        memcpy(temp, &timestamp, 4);
        memcpy(FLV_TAG_HEADER + 4, temp + 1, 3);
        FLV_TAG_HEADER[7] = temp[0];

        //处理IDRPrefixLength
        int pFrameLength = htonl(curNaluLength + 1);
        memcpy(PFRAME_PREFIX + 5, &pFrameLength, 4);

        if(!Tools::WriteFile(file, clientFd, FLV_TAG_HEADER, FLV_TAG_HEADER_LENGTH, "写入FlvHeader时发生错误。")) return false;
        if(!Tools::WriteFile(file, clientFd, PFRAME_PREFIX, PFRAME_PREFIX_LENGTH, "写入PFramePrefix时发生错误。")) return false;
    }

    int payloadLength = sendNode->endIndex - sendNode->startIndex;
    if(payloadLength > 0)
    {
        if(!Tools::WriteFile(file, clientFd, recBuffer + sendNode->startIndex, payloadLength, "写入nalu载荷时发生错误。")) return false;
    }
    else
    {
        payloadLength = REC_BUFFER_LENGTH_FLV - sendNode->startIndex;
        if(!Tools::WriteFile(file, clientFd, recBuffer + sendNode->startIndex, payloadLength, "写入nalu载荷时发生错误。")) return false;

        payloadLength = sendNode->endIndex;
        if(!Tools::WriteFile(file, clientFd, recBuffer, payloadLength, "写入nalu载荷时发生错误。")) return false;
    }

    if(sendNode->packageType == 0 || sendNode->packageType == 2)
    {
        int previousframeLength = htonl(PFRAME_PREFIX_LENGTH + curNaluLength + FLV_TAG_HEADER_LENGTH);
        if(!Tools::WriteFile(file, clientFd, &previousframeLength, 4, "写入上一帧长度时发生错误。")) return false;
    }

    return true;
}

void FlvWriter::NaluHandler()
{
    std::cout << "--------------FlvWriter线程已开始-------------" << std::endl;

    while (GetActive())
    {
        bool stopFlag = false;
        bool chaseFlag = false;
        curNaluLength = sendNode->naluLength;
        if(curNaluLength == 0)
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            continue;
        } 

        //如果当前帧的索引小于最新的IDR的索引，开始追帧
        if(sendNode->nodeIndex < GetLastestIDR())
        {
            std::cout << "跳过一条P帧" << std::endl;
            chaseFlag = true;
        }

        while (sendNode->next)
        {
            if(sendNode->naluType == 7) stopFlag = !SPSHandler();
            else if(sendNode->naluType == 8) stopFlag = !PPSHandler();
            else if(sendNode->naluType == 5) stopFlag = !IDRHandler();
            else if(sendNode->naluType == 1)
            {
                //追帧开启时跳过P帧
                if(!chaseFlag) stopFlag = !PFrameHandler();
            }
            else
            {
                std::cout << "无效的nalu类型：---------" << sendNode->naluType << "---------" << std::endl << "流转发结束" << std::endl;
                stopFlag = true;
            }

            if(stopFlag)
            {
                std::cout << "--------------FlvWriter线程已退出-------------" << std::endl; 
                SetActive(false);
                return;
            }

            DataNode* temp = sendNode;
            int type = temp->packageType;
            delete temp;
            sendNode = sendNode->next;
            if(type == 0 || type == 2) break;
        }
    }

    std::cout << "--------------FlvWriter线程已退出-------------" << std::endl;
}

bool FlvWriter::AfterGetSPSAndPPS()
{
    int index = 32;
    char profile_idc = H264SPS[1];
    Tools::SPSDecodeU(H264SPS, index);
    int chroma_array_type = 1;
    if(profile_idc == 100 || profile_idc == 110 || profile_idc == 122 || profile_idc == 244 || profile_idc == 44 || profile_idc == 83 ||
        profile_idc == 86 || profile_idc == 118 || profile_idc == 128 || profile_idc == 138|| profile_idc == 139 || profile_idc == 134 || profile_idc == 135)
    {
        int chroma_format_idc = Tools::SPSDecodeU(H264SPS, index);
        chroma_array_type = chroma_format_idc;
        if(chroma_format_idc == 3)
        {
            int separate_colour_plane_flag = ((*(H264SPS + index / 8)) >> (7 - index % 8)) & 0x01;
            if(separate_colour_plane_flag) chroma_array_type = 0;
            index++;
        }
        Tools::SPSDecodeU(H264SPS, index);
        Tools::SPSDecodeU(H264SPS, index);
        index++;
        int seq_scaling_matrix_present_flag = ((*(H264SPS + index / 8)) >> (7 - index % 8)) & 0x01;
        index++;
        if(seq_scaling_matrix_present_flag)
        {
            for(int i = 0; i < ((chroma_format_idc != 3) ? 8 : 12 ); i++) index++;
        }
    }

    Tools::SPSDecodeU(H264SPS, index);
    int pic_order_cnt_type = Tools::SPSDecodeU(H264SPS, index);
    if(pic_order_cnt_type == 0)
    {
        Tools::SPSDecodeU(H264SPS, index);
    }
    else if(pic_order_cnt_type == 1)
    {
        index++;
        Tools::SPSDecode(H264SPS, index);
        Tools::SPSDecode(H264SPS, index);
        int num_ref_frames_in_pic_order_cnt_cycle = Tools::SPSDecodeU(H264SPS, index);
        for(int i = 0; i < num_ref_frames_in_pic_order_cnt_cycle; i++) Tools::SPSDecode(H264SPS, index);
    }
    Tools::SPSDecodeU(H264SPS, index);
    index++;
    int pic_width_in_mbs_minus1 = Tools::SPSDecodeU(H264SPS, index);
    int pic_height_in_map_units_minus1 = Tools::SPSDecodeU(H264SPS, index);

    int frame_mbs_only_flag = ((*(H264SPS + index / 8)) >> (7 - index % 8)) & 0x01;
    index++;
    if(!frame_mbs_only_flag)
    {
        std::cout << "SPS的frame_mbs_only_flag为0，可能存在场编码" << std::endl;
        index++;
    }
    index++;

    int subWidth = 1;
    int subHeight = 1;
    if(chroma_array_type == 1)
    {
        subWidth = 2;
        subHeight = 2;
    }
    else if(chroma_array_type == 2) subWidth = 2;

    int frame_cropping_flag = ((*(H264SPS + index / 8)) >> (7 - index % 8)) & 0x01;
    index++;
    int frame_crop_left_offset = 0;
    int frame_crop_right_offset = 0;
    int frame_crop_top_offset = 0;
    int frame_crop_bottom_offset = 0;
    int crop_unit_x = 1;
    int crop_unit_y = 1;
    if(frame_cropping_flag)
    {
        frame_crop_left_offset = Tools::SPSDecodeU(H264SPS, index);
        frame_crop_right_offset = Tools::SPSDecodeU(H264SPS, index);
        frame_crop_top_offset = Tools::SPSDecodeU(H264SPS, index);
        frame_crop_bottom_offset = Tools::SPSDecodeU(H264SPS, index);

        if(chroma_array_type == 0) crop_unit_y = 2 - frame_mbs_only_flag;
        else
        {
            crop_unit_x = subWidth;
            crop_unit_y = subHeight * (2 - frame_mbs_only_flag);
        }
    }

    double width = (pic_width_in_mbs_minus1 + 1) * 16 - crop_unit_x * (frame_crop_left_offset + frame_crop_right_offset);
    double height = (2 - frame_mbs_only_flag) * (pic_height_in_map_units_minus1 + 1) * 16 - crop_unit_y * (frame_crop_top_offset + frame_crop_bottom_offset);

    double framRateD = frameRate;
    std::cout << "分辨率：" << width << "X" << height << "(" << CHROMA_FORMAT[chroma_array_type] << ")，帧率为：" << frameRate << std::endl;

    Tools::htonll(width, SCRIPT_TAG + FLV_SCRIPT_TAG_WIDTH_OFFSET);
    Tools::htonll(height, SCRIPT_TAG + FLV_SCRIPT_TAG_HEIRHT_OFFSET);
    Tools::htonll(framRateD, SCRIPT_TAG + FLV_SCRIPT_TAG_FRAMERATE_OFFSET);

    if(!Tools::WriteFile(file, clientFd, flvPre, FLV_HEADER_LENGTH, "写入flv文件头时发生错误。")) return false;
    if(!Tools::WriteFile(file, clientFd, SCRIPT_TAG, FLV_SCRIPT_TAG_LENGTH, "写入ScriptTag时发生错误。")) return false;

    //写入AVCSeq
    int AVCSeqLength = 16 + h264SPSLength + h264PPSLength;
    AVCSeqFrameLength = FLV_TAG_HEADER_LENGTH + AVCSeqLength + 4;
    int _AVCSeqLength = htonl(AVCSeqLength);
    memcpy(AVCSeq, &_AVCSeqLength, 4);
    AVCSeq[0] = 0x09;
    memcpy(AVCSeq + FLV_TAG_HEADER_LENGTH + 6, H264SPS + 1, 3);
    short _SPSLength = htons(h264SPSLength);
    memcpy(AVCSeq + FLV_TAG_HEADER_LENGTH + 11, &_SPSLength, 2);
    memcpy(AVCSeq + FLV_TAG_HEADER_LENGTH + 13, H264SPS, h264SPSLength);
    uint8_t PPSPre[3] = {0x01, 0x00, 0x00};
    short _PPSLength = htons(h264PPSLength);
    memcpy(PPSPre + 1, &_PPSLength, 2);
    memcpy(AVCSeq + FLV_TAG_HEADER_LENGTH + 13 + h264SPSLength, PPSPre, 3);
    memcpy(AVCSeq + FLV_TAG_HEADER_LENGTH + 16 + h264SPSLength, H264PPS, h264PPSLength);
    int _preTagLength = htonl(AVCSeqLength + FLV_TAG_HEADER_LENGTH);
    memcpy(AVCSeq + FLV_TAG_HEADER_LENGTH + AVCSeqLength, &_preTagLength, 4);

    if(!Tools::WriteFile(file, clientFd, AVCSeq, AVCSeqFrameLength, "写入AVCSeq时发生错误。")) return false;
    
    //生成IDR前缀
    IDRPrefixLength = 18 + h264SPSLength + h264PPSLength;
    uint8_t pre1[9] = {0x17, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
    memcpy(IDRPrefix, pre1, 9);
    
    int __SPSLength = htonl(h264SPSLength);
    memcpy(IDRPrefix + 5, &__SPSLength, 4);
    memcpy(IDRPrefix + 9, H264SPS, h264SPSLength);
    int __PPSLength = htonl(h264PPSLength);
    memcpy(IDRPrefix + 9 + h264SPSLength, &__PPSLength, 4);
    memcpy(IDRPrefix + 13 + h264SPSLength, H264PPS, h264PPSLength);
    uint8_t pre2[5] = {0x00, 0x00, 0x00, 0x00, 0x65};
    memcpy(IDRPrefix + 13 + h264SPSLength + h264PPSLength, pre2, 5);

    return true;
}

void FlvWriter::FlvThread(FlvWriter* flvWriter)
{
    flvWriter->NaluHandler();
}

FlvWriter::FlvWriter(char* _recBuffer, float _frameRate, int _flvClientSocketFd, DataNode* _sendNode)
    : recBuffer(_recBuffer), frameRate(_frameRate), clientFd(_flvClientSocketFd), sendNode(_sendNode)
{
    file = NULL;
    // file = fopen("test.flv", "w");
 
    curNaluLength = 0;

    hasScriptTag = false;
    h264SPSLength = 0;
    h264PPSLength = 0;
    IDRPrefixLength = 0;
    AVCSeqFrameLength = 0;

    lastestIDRIndex = -1;
}

FlvWriter::~FlvWriter()
{
    if (file != NULL) fclose(file);

    int count = 0;
    while (sendNode)
    {
        DataNode* temp = sendNode;
        sendNode = sendNode->next;
        delete temp;
        count++;
    }

    std::cout << "FlvWriter已释放，删除" << count << "个未使用数据节点" << std::endl;
}

DataNode* FlvWriter::GetSendNode()
{
    return sendNode;
}

void FlvWriter::SetLastestIDR(int index)
{
    mtxIDRIndex.lock();
    lastestIDRIndex = index;
    mtxIDRIndex.unlock();
}

int FlvWriter::GetLastestIDR()
{
    mtxIDRIndex.lock();
    int ret = lastestIDRIndex;
    mtxIDRIndex.unlock();
    return ret;
}
   
void FlvWriter::SetActive(bool _active)
{
    mtxActive.lock();
    active = _active;
    mtxActive.unlock();
}

bool FlvWriter::GetActive()
{
    mtxActive.lock();
    bool ret = active;
    mtxActive.unlock();
    return ret;
}